package com.stress.testing.plugin.apache.rocketmq.hook;

import com.stress.testing.common.TraceContext;
import com.stress.testing.common.TraceContextHolder;
import com.stress.testing.common.conf.AgentConfig;
import com.stress.testing.common.conf.PluginConfig;
import org.apache.rocketmq.client.hook.ConsumeMessageContext;
import org.apache.rocketmq.client.hook.ConsumeMessageHook;
import org.apache.rocketmq.common.MixAll;
import org.apache.rocketmq.common.message.MessageExt;

import java.util.Map;

import static com.stress.testing.common.conf.PluginConfigKey.DEFAULT_PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX;
import static com.stress.testing.common.conf.PluginConfigKey.PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX_KEY;

/**
 * 已废弃。改为通过listener包装类实现。
 *
 * @see com.stress.testing.plugin.apache.rocketmq.listener.MessageListenerOrderlyWrapper
 * @see com.stress.testing.plugin.apache.rocketmq.listener.MessageListenerConcurrentlyWrapper
 */
@Deprecated
public class PluginConsumeMessageHook implements ConsumeMessageHook {
    private AgentConfig agentConfig;
    private PluginConfig pluginConfig;

    public PluginConsumeMessageHook(AgentConfig agentConfig, PluginConfig pluginConfig) {
        this.agentConfig = agentConfig;
        this.pluginConfig = pluginConfig;
    }

    @Override
    public String hookName() {
        return PluginConsumeMessageHook.class.getSimpleName();
    }

    /**
     * 当消息的数量只有一条的时候，设置当前线程的TraceContext
     * 这个操作的意义不大
     * 本质上是想针对单条msg，设置TraceContext。由于是批量消费，无法在字节码增强阶段，确认究竟是哪条msg在消费
     * 解决方案：
     * 1、业务代码侵入，最不友好
     * 2、sdk支持，类似于maihaoche的jar，可以知道单条msg消费的起点和终点。在起点设置TraceContext，在终点移除TraceContext
     *
     * @param context
     */
    @Override
    public void consumeMessageBefore(ConsumeMessageContext context) {
        TraceContextHolder.removeTraceContext();

        if (context == null || context.getMsgList() == null || context.getMsgList().size() != 1) {
            return;
        }

        MessageExt messageExt = context.getMsgList().get(0);
        String topic = messageExt.getTopic();
        String topicPrefix = pluginConfig.get(PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX_KEY,
                DEFAULT_PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX);

        if (!topic.startsWith(topicPrefix) &&
                !topic.startsWith(MixAll.DLQ_GROUP_TOPIC_PREFIX + topicPrefix) &&
                !topic.startsWith(MixAll.RETRY_GROUP_TOPIC_PREFIX + topicPrefix)) {
            return;
        }

        TraceContext traceContext = new TraceContext(agentConfig.getTracePropertyPrefix());
        Map<String, String> properties = messageExt.getProperties();
        for (Map.Entry<String, String> entry : properties.entrySet()) {
            if (entry.getKey().startsWith(agentConfig.getTracePropertyPrefix())) {
                traceContext.setProperty(entry.getKey(), entry.getValue());
            }
        }

        // 可以简单的认为，如果msg的property中不包含trace_id,那么就不包含trace信息
        if (!traceContext.hasTraceID()) {
            return;
        }

        TraceContextHolder.setTraceContext(traceContext);
    }

    @Override
    public void consumeMessageAfter(ConsumeMessageContext context) {
        TraceContextHolder.removeTraceContext();
    }
}
